{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Issue [#99](https://github.com/bsorrentino/langgraph4j/issues/99) by [zu1k](https://github.com/zu1k)\n",
    "\n",
    "Verify \"**Edge Update Logic Error During Subgraph Processing**\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "var userHomeDir = System.getProperty(\"user.home\");\n",
    "var localRespoUrl = \"file://\" + userHomeDir + \"/.m2/repository/\";\n",
    "var langchain4jVersion = \"1.0.1\";\n",
    "var langchain4jbeta = \"1.0.1-beta6\";\n",
    "var langgraph4jVersion = \"1.6-SNAPSHOT\";"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash \n",
    "rm -rf \\{userHomeDir}/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/bsc/langgraph4j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%dependency /add-repo local \\{localRespoUrl} release|never snapshot|always\n",
    "// %dependency /list-repos\n",
    "%dependency /add org.slf4j:slf4j-jdk14:2.0.9\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-core:\\{langgraph4jVersion}\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-langchain4j:\\{langgraph4jVersion}\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-agent-executor:\\{langgraph4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j:\\{langchain4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j-open-ai:\\{langchain4jVersion}\n",
    "\n",
    "%dependency /resolve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Initialize Logger**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "try( var file = new java.io.FileInputStream(\"./logging.properties\")) {\n",
    "    java.util.logging.LogManager.getLogManager().readConfiguration( file );\n",
    "}\n",
    "\n",
    "var log = org.slf4j.LoggerFactory.getLogger(\"AdaptiveRag\");\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## State declaration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import org.bsc.langgraph4j.action.NodeAction;\n",
    "import org.bsc.langgraph4j.state.AgentState;\n",
    "import org.bsc.langgraph4j.StateGraph;\n",
    "\n",
    "import static org.bsc.langgraph4j.StateGraph.END;\n",
    "import static org.bsc.langgraph4j.StateGraph.START;\n",
    "import static org.bsc.langgraph4j.action.AsyncEdgeAction.edge_async;\n",
    "import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async;\n",
    "\n",
    "import java.util.Map;\n",
    "import java.util.Optional;\n",
    "\n",
    "class State extends AgentState {\n",
    "\n",
    "    public State(Map<String, Object> initData) {\n",
    "        super(initData);\n",
    "    }\n",
    "\n",
    "    public Optional<String> intent() {\n",
    "        return  value(\"intent\");\n",
    "    }\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Subgraph definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "// SubGraph1 Definition\n",
    "StateGraph<State> subGraph1 = new StateGraph<>( State::new )\n",
    "            .addNode(\"work\", node_async(state -> Map.of(\"step\", \"work1\")))\n",
    "            .addEdge(START, \"work\")\n",
    "            .addEdge(\"work\", END)\n",
    "            ;\n",
    "\n",
    "// SubGraph2 Definition\n",
    "StateGraph<State> subGraph2 = new StateGraph<>( State::new )\n",
    "            .addNode(\"work\", node_async(state -> Map.of(\"step\", \"work2\")))\n",
    "            .addEdge(START, \"work\")\n",
    "            .addEdge(\"work\", END)\n",
    "            ;\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IntentRecognize node declaration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class IntentRecognizeNode implements NodeAction<State> {\n",
    "\n",
    "    String intent;\n",
    "\n",
    "    public void setIntent( String intent ) {\n",
    "        this.intent = intent;\n",
    "    }\n",
    "\n",
    "    public String getIntent() {\n",
    "        return intent;\n",
    "    }\n",
    "\n",
    "    @Override\n",
    "    public Map<String, Object> apply(State state) {\n",
    "        return Map.of( \"intent\", intent );\n",
    "    }\n",
    "\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Subgraph as StateGraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "EXPLAIN \n",
      "START \n",
      "NodeOutput{node=__START__, state={input=explain me}} \n",
      "NodeOutput{node=intent_recognize, state={input=explain me, intent=explain}} \n",
      "NodeOutput{node=subAgent1-work, state={input=explain me, intent=explain, step=work1}} \n",
      "NodeOutput{node=__END__, state={input=explain me, intent=explain, step=work1}} \n",
      "QUERY \n",
      "START \n",
      "NodeOutput{node=__START__, state={input=search for}} \n",
      "NodeOutput{node=intent_recognize, state={input=search for, intent=query}} \n",
      "NodeOutput{node=subAgent2-work, state={input=search for, intent=query, step=work2}} \n",
      "NodeOutput{node=__END__, state={input=search for, intent=query, step=work2}} \n"
     ]
    }
   ],
   "source": [
    "var intentRecognizeNode = new IntentRecognizeNode();\n",
    "\n",
    "// MainGraph Definition\n",
    "var graph = new StateGraph<>( State::new )\n",
    "        .addNode(\"intent_recognize\", node_async(intentRecognizeNode))\n",
    "        .addNode(\"subAgent1\", subGraph1 )\n",
    "        .addNode(\"subAgent2\", subGraph2 )\n",
    "        .addEdge(START, \"intent_recognize\")\n",
    "        .addConditionalEdges(\"intent_recognize\",\n",
    "                edge_async( state ->\n",
    "                        state.intent().orElseThrow() ),\n",
    "                Map.of(\"explain\", \"subAgent1\",\n",
    "                        \"query\", \"subAgent2\"\n",
    "                )\n",
    "        )\n",
    "        .addEdge(\"subAgent1\", END)\n",
    "        .addEdge(\"subAgent2\", END)\n",
    "        ;\n",
    "\n",
    "var workflow = graph.compile();\n",
    "\n",
    "// System.out.println( workflow.getGraph( GraphRepresentation.Type.PLANTUML, \"\", false ));\n",
    "\n",
    "// EXPLAIN\n",
    "log.info( \"EXPLAIN\");\n",
    "intentRecognizeNode.setIntent(\"explain\");\n",
    "for( var output : workflow.stream( Map.of(\"input\", \"explain me\") ) ) {\n",
    "       log.info( \"{}\", output );\n",
    "}\n",
    "\n",
    "// QUERY\n",
    "log.info( \"QUERY\");\n",
    "intentRecognizeNode.setIntent(\"query\");\n",
    "for( var output : workflow.stream( Map.of(\"input\", \"search for\") ) ) {\n",
    "        log.info( \"{}\", output );\n",
    "}\n",
    "         \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Subgraph as CompiledGraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "EXPLAIN \n",
      "START \n",
      "NodeOutput{node=__START__, state={input=explain me}} \n",
      "START \n",
      "NodeOutput{node=intent_recognize, state={input=explain me, intent=explain}} \n",
      "NodeOutput{node=__START__, state={input=explain me, intent=explain}} \n",
      "NodeOutput{node=work, state={input=explain me, intent=explain, step=work1}} \n",
      "NodeOutput{node=__END__, state={input=explain me, intent=explain, step=work1}} \n",
      "NodeOutput{node=subAgent1, state={input=explain me, intent=explain, step=work1}} \n",
      "NodeOutput{node=__END__, state={input=explain me, intent=explain, step=work1}} \n",
      "QUERY \n",
      "START \n",
      "NodeOutput{node=__START__, state={input=search for}} \n",
      "START \n",
      "NodeOutput{node=intent_recognize, state={input=search for, intent=query}} \n",
      "NodeOutput{node=__START__, state={input=search for, intent=query}} \n",
      "NodeOutput{node=work, state={input=search for, intent=query, step=work2}} \n",
      "NodeOutput{node=__END__, state={input=search for, intent=query, step=work2}} \n",
      "NodeOutput{node=subAgent2, state={input=search for, intent=query, step=work2}} \n",
      "NodeOutput{node=__END__, state={input=search for, intent=query, step=work2}} \n"
     ]
    }
   ],
   "source": [
    "var intentRecognizeNode = new IntentRecognizeNode();\n",
    "\n",
    "// MainGraph Definition\n",
    "var graph = new StateGraph<>( State::new )\n",
    "        .addNode(\"intent_recognize\", node_async(intentRecognizeNode))\n",
    "        .addNode(\"subAgent1\", subGraph1.compile() )\n",
    "        .addNode(\"subAgent2\", subGraph2.compile() )\n",
    "        .addEdge(START, \"intent_recognize\")\n",
    "        .addConditionalEdges(\"intent_recognize\",\n",
    "                edge_async( state ->\n",
    "                        state.intent().orElseThrow() ),\n",
    "                Map.of(\"explain\", \"subAgent1\",\n",
    "                        \"query\", \"subAgent2\"\n",
    "                )\n",
    "        )\n",
    "        .addEdge(\"subAgent1\", END)\n",
    "        .addEdge(\"subAgent2\", END)\n",
    "        ;\n",
    "\n",
    "var workflow = graph.compile();\n",
    "\n",
    "// System.out.println( workflow.getGraph( GraphRepresentation.Type.PLANTUML, \"\", false ));\n",
    "\n",
    "// EXPLAIN\n",
    "log.info( \"EXPLAIN\");\n",
    "intentRecognizeNode.setIntent(\"explain\");\n",
    "for( var output : workflow.stream( Map.of(\"input\", \"explain me\") ) ) {\n",
    "       log.info( \"{}\", output );\n",
    "}\n",
    "\n",
    "// QUERY\n",
    "log.info( \"QUERY\");\n",
    "intentRecognizeNode.setIntent(\"query\");\n",
    "for( var output : workflow.stream( Map.of(\"input\", \"search for\") ) ) {\n",
    "        log.info( \"{}\", output );\n",
    "}\n",
    "         \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Java (rjk 2.2.0)",
   "language": "java",
   "name": "rapaio-jupyter-kernel"
  },
  "language_info": {
   "codemirror_mode": "java",
   "file_extension": ".jshell",
   "mimetype": "text/x-java-source",
   "name": "java",
   "nbconvert_exporter": "script",
   "pygments_lexer": "java",
   "version": "22.0.2+9-70"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
